from typing import List, Any
from ortools.linear_solver import pywraplp as lp


class AreaSolver:
    IDLE = 0
    OPTIMAL = 1
    FEASIBLE = 2
    INFEASIBLE = 3

    def __init__(self,
            unit,                           # 排班单元
            staff_ids,                       # 当天此班型出勤的员工数
            ):
        self._unit = unit
        self._staff_ids = staff_ids

        all_staff_ids = [staff.id for staff in unit.staff_list]
        self._staff_indexes = [all_staff_ids.index(s_id) for s_id in staff_ids]
        self._staff_cnt = len(staff_ids)

        self._area_data = unit.area_data    # 当天此班型时段内区域航班分布
        self._area_summary = [item["area_summary"] for item in unit.exist_summary] # 当月已有数据，所有人员工作区域分布
        self._status = AreaSolver.IDLE
        
        self._area_list = []
        self._area_flight_cnt = []
        for key, value in unit.area_data.items():
            self._area_list.append(key)
            self._area_flight_cnt.append(len(value))
        self._flight_cnt = sum(self._area_flight_cnt)
        self._area_cnt = len(self._area_list)

        # 超参数
        self._max_deviation = 2

    def solve(self):
        # 计算每个区域需要的人数
        area_staffcnt_list = [0 for j in range(self._area_cnt)]
        for j in range(self._area_cnt):
            area_staffcnt_list[j] = self._area_flight_cnt[j] / self._flight_cnt * self._staff_cnt

        # 参数
        deviation = 0
        optimal_flag = False
        EXTREME_VAL = max([max(staff_area_data) for staff_area_data in self._area_summary]) + 1

        while (not optimal_flag) and (deviation <= self._max_deviation):
            deviation += 1

            # 初始化求解器
            solver = lp.Solver("AreaSolver", lp.Solver.CBC_MIXED_INTEGER_PROGRAMMING)

            # 变量
            var_matrix = [[solver.IntVar(0, 1, f"Staff{i} Area{j}") 
                    for j in range(self._area_cnt)] 
                    for i in range(self._staff_cnt)]
            extreme_var_matrix = [
                        {"max": solver.IntVar(0, EXTREME_VAL, f"Staff{i} Worktime Cnt Max"),
                         "min": solver.IntVar(0, EXTREME_VAL, f"Staff{i} Worktime Cnt Min")}
                    for i in range(self._staff_cnt)]

            # 约束
            # 区域员工数量约束
            for j in range(self._area_cnt):
                solver.Add(solver.Sum(
                    [var_matrix[i][j] for i in range(self._staff_cnt)])
                    >= area_staffcnt_list[j] - deviation
                )
                solver.Add(solver.Sum(
                    [var_matrix[i][j] for i in range(self._staff_cnt)])
                    <= area_staffcnt_list[j] + deviation
                )

            # 出勤区域唯一性约束
            for i in range(self._staff_cnt):
                solver.Add(solver.Sum(
                    [var_matrix[i][j] for j in range(self._area_cnt)])
                    == 1
                )

            # 区域均衡性极值变量的约束
            for i in range(self._staff_cnt):
                for j in range(self._area_cnt):
                    summary = self._area_summary[self._staff_indexes[i]]
                    solver.Add(extreme_var_matrix[i]["min"] <= summary[j] + var_matrix[i][j])
                    solver.Add(extreme_var_matrix[i]["max"] >= summary[j] + var_matrix[i][j])


            # 目标
            obj = solver.Objective()
            for i in range(self._staff_cnt):
                obj.SetCoefficient(extreme_var_matrix[i]["max"], 1)
                obj.SetCoefficient(extreme_var_matrix[i]["min"], -1)
                # for j in range(self._area_cnt):
                #     obj.SetCoefficient(var_matrix[i][j], 1)
            obj.SetMinimization()

            # 求解
            solver.EnableOutput()
            status = solver.Solve()
            if status == lp.Solver.OPTIMAL:
                self._status = AreaSolver.OPTIMAL
                optimal_flag = True

                # 计算结果数据
                staff_area_list = []
                for i in range(self._staff_cnt):
                    for j in range(self._area_cnt):  
                        if var_matrix[i][j].solution_value() == 1:
                            staff_area_list.append((self._staff_ids[i], self._area_list[j]))
                            break
                self._res = staff_area_list

                return self._status

        self._status = AreaSolver.INFEASIBLE
        return self._status

    def calc_return_data(self):
        return self._res


# 先分组后分区域的算法

# class AreaSolver:
#     IDLE = 0
#     OPTIMAL = 1
#     FEASIBLE = 2
#     INFEASIBLE = 3

#     def __init__(self,
#             group_data: List[List[Any]],    # 分组信息List，其中每个 item 是本组的员工ID List
#             unit,
#             ):
#         self._group_data = group_data
#         self._area_data = unit.area_data
#         self._status = AreaSolver.IDLE
        
#         group_staff_cnts = [len(group) for group in group_data]
#         self._max_group_staff_cnt = max(group_staff_cnts)
#         self._min_group_staff_cnt = min(group_staff_cnts)
#         self._staff_cnt = sum(group_staff_cnts)
#         self._group_cnt = len(group_data)

#         self._area_list = []
#         self._area_flight_cnt = []
#         for key, value in self._area_data.items():
#             self._area_list.append(key)
#             self._area_flight_cnt.append(len(value))
#         self._flight_cnt = sum(self._area_flight_cnt)
#         self._area_cnt = len(self._area_list)

#     def solve(self):
#         # 计算每个区域需要的人数
#         area_staffcnt_list = [0 for j in range(len(self._area_list))]
#         for j in range(len(self._area_list)):
#             area_staffcnt_list[j] = self._area_flight_cnt[j] / self._flight_cnt * self._staff_cnt

#         # 参数
#         deviation = 0
#         optimal_flag = False

#         while (not optimal_flag) and (deviation <= self._max_group_staff_cnt):
#             deviation += 1

#             # 初始化求解器
#             solver = lp.Solver("AreaSolver", lp.Solver.CBC_MIXED_INTEGER_PROGRAMMING)

#             # 变量
#             var_matrix = [[solver.IntVar(0, 1, f"Group{i} Area{j}") 
#                     for j in range(self._area_cnt)] 
#                     for i in range(self._group_cnt)]

#             # 约束
#             # 区域员工数量约束
#             for j in range(self._area_cnt):
#                 solver.Add(solver.Sum(
#                     [var_matrix[i][j] * len(self._group_data[i]) for i in range(self._group_cnt)])
#                     >= area_staffcnt_list[j] - deviation
#                 )
#                 solver.Add(solver.Sum(
#                     [var_matrix[i][j] * len(self._group_data[i]) for i in range(self._group_cnt)])
#                     <= area_staffcnt_list[j] + deviation
#                 )

#             for i in range(self._group_cnt):
#                 solver.Add(solver.Sum(
#                     [var_matrix[i][j] for j in range(self._area_cnt)])
#                     == 1
#                 )

#             # 目标
#             obj = solver.Objective()
#             for i in range(self._group_cnt):
#                 for j in range(self._area_cnt):
#                     obj.SetCoefficient(var_matrix[i][j], 1)
#             obj.SetMinimization()

#             # 求解
#             solver.EnableOutput()
#             status = solver.Solve()
#             if status == lp.Solver.OPTIMAL:
#                 self._status = AreaSolver.OPTIMAL
#                 optimal_flag = True

#                 # 计算结果数据
#                 staff_area_list = []
#                 for i in range(self._group_cnt):
#                     for j in range(self._area_cnt):  
#                         if var_matrix[i][j].solution_value() == 1:
#                             for staff_id in self._group_data[i]:
#                                 staff_area_list.append((staff_id, self._area_list[j]))
#                             break
#                 self._res = staff_area_list

#                 return self._status

#         self._status = AreaSolver.INFEASIBLE
#         return self._status

#     def calc_return_data(self):
#         return self._res
